"""
基于联邦学习的跨市场套利策略
策略特点：
1. 使用联邦学习框架聚合多市场特征
2. 实现隐私保护的套利信号发现
3. 支持动态权重调整
"""

import torch
import torch.nn as nn
import numpy as np
from torch.utils.data import Dataset, DataLoader

class FederatedModel(nn.Module):
    def __init__(self, input_dim=10, hidden_dim=64):
        super().__init__()
        self.shared_encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        self.market_heads = nn.ModuleDict({
            'us': nn.Linear(hidden_dim, 1),
            'hk': nn.Linear(hidden_dim, 1),
            'cn': nn.Linear(hidden_dim, 1)
        })
        
    def forward(self, x, market):
        features = self.shared_encoder(x)
        return self.market_heads[market](features)

class ArbitrageStrategy:
    def __init__(self):
        self.global_model = FederatedModel()
        self.local_models = {
            'us': FederatedModel(),
            'hk': FederatedModel(),
            'cn': FederatedModel()
        }
        self.optimizers = {
            m: torch.optim.Adam(model.parameters())
            for m, model in self.local_models.items()
        }
        
    def federated_train(self, dataloaders, rounds=10):
        for _ in range(rounds):
            # 本地训练
            local_updates = {}
            for market in self.local_models:
                self._local_train(market, dataloaders[market])
                local_updates[market] = self._get_model_diff(market)
            
            # 联邦聚合
            self._aggregate_updates(local_updates)
            
    def _local_train(self, market, dataloader, epochs=3):
        model = self.local_models[market]
        model.train()
        for _ in range(epochs):
            for x, y in dataloader:
                self.optimizers[market].zero_grad()
                pred = model(x, market)
                loss = nn.MSELoss()(pred, y)
                loss.backward()
                self.optimizers[market].step()
                
    def _get_model_diff(self, market):
        """获取本地模型与全局模型的差异"""
        return {
            name: param.data - self.global_model.state_dict()[name]
            for name, param in self.local_models[market].named_parameters()
        }
        
    def _aggregate_updates(self, updates):
        """聚合各市场更新"""
        avg_update = {
            name: torch.stack([updates[m][name] for m in updates]).mean(0)
            for name in updates['us']
        }
        
        # 更新全局模型
        global_state = self.global_model.state_dict()
        for name in avg_update:
            global_state[name] += avg_update[name]
        self.global_model.load_state_dict(global_state)
        
        # 更新本地模型
        for market in self.local_models:
            local_state = self.local_models[market].state_dict()
            for name in avg_update:
                local_state[name] = global_state[name]
            self.local_models[market].load_state_dict(local_state)
            
    def detect_arbitrage(self, market1, market2, features1, features2):
        """检测套利机会"""
        with torch.no_grad():
            spread = (self.global_model(features1, market1) - 
                     self.global_model(features2, market2))
        return spread.item()

if __name__ == '__main__':
    # 示例用法
    np.random.seed(42)
    
    # 模拟多市场数据
    us_data = np.random.randn(100, 10)
    hk_data = np.random.randn(100, 10)
    cn_data = np.random.randn(100, 10)
    
    # 创建数据加载器
    class MarketDataset(Dataset):
        def __init__(self, data):
            self.x = torch.FloatTensor(data[:, :-1])
            self.y = torch.FloatTensor(data[:, -1:])
            
        def __len__(self):
            return len(self.x)
            
        def __getitem__(self, idx):
            return self.x[idx], self.y[idx]
    
    dataloaders = {
        'us': DataLoader(MarketDataset(us_data), batch_size=32),
        'hk': DataLoader(MarketDataset(hk_data), batch_size=32),
        'cn': DataLoader(MarketDataset(cn_data), batch_size=32)
    }
    
    # 训练联邦模型
    strategy = ArbitrageStrategy()
    strategy.federated_train(dataloaders)
    
    # 检测套利机会
    test_us = torch.FloatTensor(np.random.randn(10))
    test_hk = torch.FloatTensor(np.random.randn(10))
    spread = strategy.detect_arbitrage('us', 'hk', test_us, test_hk)
    print(f"市场价差: {spread:.4f}")